#include "RSAMessage.h"
#include <cstdlib>
#include <cstdio>
#include <cstring>
#include <openssl/rsa.h>

namespace crypto
{
  namespace rsa
  {
    RSAMessage::RSAMessage()
    {
      // reseting all pointers
      this->input = NULL;
      this->inputLength = 0;
      this->output = NULL;
      this->outputLength = 0;
    }

    RSAMessage::~RSAMessage() 
    {
      if(this->input != NULL) 
      {
        delete[] this->input;
      }
      if(this->output != NULL) 
      {
        delete[] this->output;
      }
    }

    void RSAMessage::setInput(char *input, int inputLength) 
    {
      // if input is not empty, reset it
      if (this->input != NULL) 
      {
        delete[] this->input;
      }

      this->input = new char[inputLength];
      if (this->input == NULL) 
      {
        fprintf(stderr, "class RSACryptMessage::setInput: "
            "cannot allocate memory!\n");
        exit(1);
      }

      this->inputLength = inputLength;
      memcpy(this->input, input, inputLength);
    }

    int RSAMessage::getOutputLength() 
    {
      if (this->output == NULL) 
      {
        return -1;
      }
      return this->outputLength;
    }

    char * RSAMessage::getOutput() 
    {
      if (this->output == NULL) 
      {
        return NULL;
      }
      return this->output;
    }

    void RSAMessage::encrypt(RSAKey *key)
    {
      if (key == NULL)
      {
        throw "wrong argument";
      }
      if (this->input == NULL) 
      {
        fprintf(stderr, "class RSACryptMessage::encryptMessage(): "
            "input message is not setted!\n");
        exit(1);
      }
      if (this->inputLength <= 0)
      {
        fprintf(stderr, "class RSACryptMessage::encryptMessage(): "
            "input message length is wrong!\n");
        exit(1);
      }

      // create output
      int maxBlockLength = key->size() - this->MESSAGE_REST;
      int blockCount = (this->inputLength + maxBlockLength - 1) / maxBlockLength;
      if (this->output != NULL)
      {
        delete[] this->output;
      }
      this->output = new char[key->size() * blockCount];
      if (this->output == NULL)
      {
        fprintf(stderr, "class RSACryptMessage::encryptMessage: "
            "cannot allocate memory!\n");
        exit(1);
      }
      int outputPointer = 0;

      // encrypt input to output by blocks
      for (int i = 0; i < blockCount; ++i) 
      {
        int messageLength;
        if (i < blockCount - 1)
        {
          messageLength = maxBlockLength;
        }
        else
        {
          messageLength = this->inputLength - i * maxBlockLength;
          printf("last messageLength: %d\n", messageLength);
        }
        outputPointer += RSA_public_encrypt(messageLength,
            (unsigned char *)(&this->input[i * maxBlockLength]),
            (unsigned char *)(&this->output[outputPointer]),
            key->toRSA(), RSA_PKCS1_OAEP_PADDING);
      }
      this->outputLength = outputPointer;
    }

    void RSAMessage::decrypt(RSAKey *key)
    {
      if (key == NULL)
      {
        throw "wrong argument!";
      }
      if (key->getType() != RSAKey::PRIVATE)
      {
        throw "wrong key type!";
      }
      if (this->input == NULL)
      {
        fprintf(stderr, "class RSACryptMessage::encryptMessage(): "
            "input message is not setted!\n");
        exit(1);
      }
      if (this->inputLength <= 0)
      {
        fprintf(stderr, "class RSACryptMessage::encryptMessage(): "
            "input message length is wrong!\n");
        exit(1);
      }

      // create output
      if (this->output != NULL)
      {
        delete[] this->output;
      }
      this->output = new char[this->inputLength];
      if (this->output == NULL)
      {
        fprintf(stderr, "class RSACryptMessage::decryptMessage: "
            "cannot allocate memory!\n");
        exit(1);
      }

      int blockLength = key->size();
      int blockCount = this->inputLength / blockLength;
      int outputPointer = 0;

      for (int i = 0; i < blockCount; ++i)
      {
        outputPointer += RSA_private_decrypt(blockLength,
            (unsigned char *)(&this->input[i * blockLength]),
            (unsigned char *)(&this->output[outputPointer]),
            key->toRSA(), RSA_PKCS1_OAEP_PADDING);
      }
      this->outputLength = outputPointer;
    }

  }
}
